docs: fix tree packing algorithm name and token alignment constraint#1205
docs: fix tree packing algorithm name and token alignment constraint#1205guozhihao-224 wants to merge 1 commit into
Conversation
1. Change "first-fit decreasing" to "first-fit". The code in _greedy_build_tries() iterates sequences in original order without sorting by decreasing size. 2. Document the full token alignment constraint: lcm(BLOCK_SIZE, parallel_size) instead of just BLOCK_SIZE (128). When parallel_size divides 128 (common case) this equals 128, but can be larger with unusual parallel configs. Ref: areal-project#1165 (Findings areal-project#6, areal-project#7) Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request updates the tree training documentation to clarify that max_tokens_per_mb must be a multiple of the least common multiple of the block size and parallel size. It also updates the description of the greedy packing strategy from first-fit decreasing to first-fit. Feedback was provided to generalize the definition of parallel_size to ensure accuracy across different backend engines like Megatron and Archon.
| `BLOCK_SIZE` (128). | ||
| NOTE: When tree training is enabled, `max_tokens_per_mb` must be a multiple of | ||
| `lcm(BLOCK_SIZE, parallel_size)`, where `BLOCK_SIZE` is 128 and | ||
| `parallel_size = tp_size × sp_size`. In most configurations where `parallel_size` |
There was a problem hiding this comment.
The definition of parallel_size as tp_size × sp_size is accurate for the FSDP backend, but the Megatron and Archon engines currently only pass the tensor parallel degree (tp_size) to the packing function. To ensure the documentation is accurate across all supported backends, consider clarifying that parallel_size represents the product of parallelism dimensions that require block alignment for the specific engine being used.
| `parallel_size = tp_size × sp_size`. In most configurations where `parallel_size` | |
| `parallel_size` (the product of parallelism dimensions requiring block alignment, e.g., `tp_size × sp_size` for FSDP). In most configurations where `parallel_size` |
1 participant
Summary
Two corrections to the tree training reference documentation:
Packing algorithm: Change "first-fit decreasing" to "first-fit". The code in
_greedy_build_tries()iterates sequences in original order without sorting by decreasing size. (A proper FFD implementation exists inareal/utils/seqpack.pyfor trajectory redistribution, but tree packing intentionally uses plain first-fit to preserve prefix locality.)Token alignment constraint: Document the full constraint
lcm(BLOCK_SIZE, parallel_size)instead of justBLOCK_SIZE(128). In most configurations whereparallel_sizedivides 128 this equals 128, but the actual code attree.py:341usesmath.lcm()which can be stricter with unusual parallel configs.Verification
_greedy_build_tries()attree.py:476-523math.lcm(BLOCK_SIZE, parallel_size)attree.py:341Test plan
pre-commit run --all-filespasses./docs/build_all.shbuilds without errorsRef: #1165 (Findings #6, #7)
🤖 Generated with Claude Code